import akshare as ak
import pandas as pd
import numpy as np
import time
import talib

# 获取所有股票代码（这里假设akshare有相应功能获取全部代码，实际可能需要调整）
def get_all_stock_codes():
    try:
        print(f"[{time.ctime()}] 获取所有股票代码...")
        # 获取 A 股股票代码和名称数据
        stock_info = ak.stock_info_a_code_name()

        def add_prefix(code):
            if code.startswith("300"):
                return "sz" + code
            elif code.startswith("60"):
                return "sh" + code
            elif code.startswith("00"):
                return "sz" + code
            else:
                return None

        # 为股票代码添加前缀并筛选
        stock_info["code_with_prefix"] = stock_info["code"].apply(add_prefix)
        filtered_stock_info = stock_info.dropna(subset=['code_with_prefix'])
        print(f"[{time.ctime()}] 完成获取所有股票代码，共 {len(filtered_stock_info)} 个。")
        return filtered_stock_info
    except Exception as e:
        print(f"获取股票代码时出错: {e}")
        return []

# 获取股票数据
def get_stock_data(symbol, start_date, end_date):
    try:
        print(f"[{time.ctime()}] 正在获取 {symbol} 的股票数据...")
        stock_data = ak.stock_zh_a_daily(symbol=symbol, start_date=start_date, end_date=end_date)
        print(f"[{time.ctime()}] 成功获取 {symbol} 的股票数据。")
        return stock_data
    except Exception as e:
        print(f"[{time.ctime()}] 获取 {symbol} 股票数据时出错: {e}")
        return pd.DataFrame()

# 交易模型函数，判断股票明天上涨的概率，并给出买入价格和卖出价格
# 交易模型函数，判断股票明天上涨的概率，并给出买入价格和卖出价格
# 交易模型函数，判断股票明天上涨的概率，并给出买入价格和卖出价格
def trading_model(simple_trading, symbol, start_date, end_date):
    data = get_stock_data(symbol, start_date, end_date)
    if data.empty:
        return 0, None, None
    # 使用talib计算技术指标
    data['MA5'] = talib.SMA(data['close'], timeperiod=5)
    data['MA10'] = talib.SMA(data['close'], timeperiod=10)
    data['RSI'] = talib.RSI(data['close'])
    upperband, middleband, lowerband = talib.BBANDS(data['close'])
    data['BOLL_UPPER'] = upperband
    data['BOLL_MA'] = middleband
    data['BOLL_LOWER'] = lowerband
    data['OBV'] = talib.OBV(data['close'], data['volume'])
    data['WR'] = talib.WILLR(data['high'], data['low'], data['close'])
    # 手动计算乖离率(BIAS)，这里以6日为例
    ma6 = talib.SMA(data['close'], timeperiod=6)
    data['BIAS'] = ((data['close'] - ma6) / ma6) * 100
    adx = talib.ADX(data['high'], data['low'], data['close'])
    plus_di = talib.PLUS_DI(data['high'], data['low'], data['close'])
    minus_di = talib.MINUS_DI(data['high'], data['low'], data['close'])
    data['ADX'] = adx
    data['PLUS_DI'] = plus_di
    data['MINUS_DI'] = minus_di
    total_count = len(data) - 1
    up_count = 0
    buy_price = None
    sell_price = None
    for i in range(len(data) - 1):
        current_rsi = data['RSI'][i]
        current_ma5 = data['MA5'][i]
        current_ma10 = data['MA10'][i]
        current_price = data['close'][i]
        current_obv = data['OBV'][i]
        current_wr = data['WR'][i]
        current_bias = data['BIAS'][i]
        current_adx = data['ADX'][i]
        plus_di = data['PLUS_DI'][i]
        minus_di = data['MINUS_DI'][i]
        next_day_price = data['close'][i + 1]

        if simple_trading:
            # 简单策略：RSI小于70，5日均线大于10日均线，价格低于布林线上轨，OBV上升，WR大于80，BIAS小于10，ADX大于25，+DI大于 - DI，明天价格上涨
            if current_rsi < 70 and current_ma5 > current_ma10 and \
                    current_price < data['BOLL_UPPER'][i] and \
                    current_obv > data['OBV'][i - 1] and \
                    current_wr > 80 and \
                    current_bias < 10 and \
                    current_adx > 25 and plus_di > minus_di and \
                    next_day_price > current_price:
                up_count += 1
                buy_price = current_price
                # 止盈为买入价的10%，止损为买入价的5%
                sell_price = current_price * 1.1 if next_day_price > current_price * 1.1 else (current_price * 0.95 if next_day_price < current_price * 0.95 else next_day_price)
    probability = up_count / total_count if total_count > 0 else 0
    return probability, buy_price, sell_price

# 示例用法
today = pd.Timestamp.now().strftime('%Y-%m-%d')
start_date = (pd.Timestamp(today) - pd.Timedelta(days=120)).strftime('%Y-%m-%d')
stock_codes = get_all_stock_codes()
for symbol in stock_codes['code_with_prefix']:
    probability, buy_price, sell_price = trading_model(simple_trading=True, symbol=symbol, start_date=start_date, end_date=today)
    print(f"股票代码: {symbol}，明天上涨的概率为: {probability}")
    print(f"买入价格: {buy_price}")
    print(f"卖出价格: {sell_price}")